"""
Copyright (c) 2025 Huawei Technologies Co., Ltd.
This file is a part of the CANN Open Software.
Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
Please refer to the License for details. You may not use this file except in compliance with the License.
THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
See LICENSE in the root of the software repository for the full text of the License.
"""

import numpy as np  # numpy：数据处理常用库（需要下载）
import argparse     # argparse：解析命令行调用脚本时的命令行输入（Python自带）
import os           # os：文件操作（Python自带）
import shutil       # shutil：文件操作（Python自带）

# gen random data in any shape
def gen_random_data_by_shape_and_type(left, right, shape, target_type):
    target = np.random.uniform(left, right, shape).astype(target_type)
    if shape == (1):
        return target[0]
    return target



# 目前只考虑了全行优先的输入
def gen_data_cal_expect_write_bin(wholeM, wholeN, wholeK, batchCount, check):
    
    print("Get input data: wholeM=%d, wholeN=%d, wholeK=%d, batchCount=%d"%(wholeM, wholeN, wholeK, batchCount))
    print("="*10)

    if check:
        print("wholeM: "+str(wholeM)
                +"\twholeN: "+str(wholeN)
                +"\twholeK: "+str(wholeK)
                +"\tbatchCount: "+str(batchCount))

    if os.path.exists("data/"):
        shutil.rmtree("data/")
    os.mkdir("data/")

    alpha = gen_random_data_by_shape_and_type(-1, 1, (1), np.float16)
    alpha.tofile("data/alpha.bin", sep= " " if check else "")
    if check:
        print("alpha: "+str(alpha))

    beta  = gen_random_data_by_shape_and_type(-1, 1, (1), np.float16)
    beta.tofile("data/beta.bin", sep=" " if check else "")
    if check:
        print("beta: "+str(beta))

    with open(f"data/Q.bin", "a" if check else "ab") as Q_file, \
         open(f"data/KT.bin", "a" if check else "ab") as KT_file, \
         open(f"data/C1.bin", "a" if check else "ab") as C1_file, \
         open(f"data/V.bin", "a" if check else "ab") as V_file, \
         open(f"data/C2.bin", "a" if check else "ab") as C2_file, \
         open(f"data/Parameters.bin", "a" if check else "ab") as Parameters_file, \
         open(f"data/C3.bin", "a" if check else "ab") as C3_file, \
         open(f"data/validM.bin", "a" if check else "ab") as validM_file, \
         open(f"data/maskA.bin", "a" if check else "ab") as maskA_file, \
         open(f"data/expect_result1.bin", "a" if check else "ab") as expect_result1_file, \
         open(f"data/expect_result2.bin", "a" if check else "ab") as expect_result2_file, \
         open(f"data/expect_result3.bin", "a" if check else "ab") as expect_result3_file:

        for batchNum in range(batchCount):

            print("BatchNum %d generating..."%(batchNum))

            # np.array([wholeM]).astype(np.int64)[0] # 
            validM = gen_random_data_by_shape_and_type(wholeM - 1, wholeM, (1), np.int64)

            if check: 
                print("validM: "+str(validM))
            
            shape_maskA_valid = (1, validM)
            
            shape_Q_valid = (validM, wholeK)
            shape_KT_valid = (wholeK, validM)
            shape_C1_valid = (validM, validM)  # 同expect_result1_valid shape

            shape_V_valid = (validM, wholeK)
            shape_C2_valid = (validM, wholeK)  # 同expect_result2_valid shape

            shape_Parameters_valid = (wholeK, wholeN)
            shape_C3_valid = (validM, wholeN)   # 同expect_result3_valid shape

            l, r = -0.125, 0.125
            
            maskA_valid = np.ones(shape_maskA_valid, np.float16)

            Q_valid = gen_random_data_by_shape_and_type(l, r, shape_Q_valid, np.float16)
            KT_valid = gen_random_data_by_shape_and_type(l, r, shape_KT_valid, np.float16)
            C1_valid = gen_random_data_by_shape_and_type(l, r, shape_C1_valid, np.float16 )

            V_valid = gen_random_data_by_shape_and_type(l, r, shape_V_valid, np.float16)
            C2_valid = gen_random_data_by_shape_and_type(l, r, shape_C2_valid, np.float16)

            Parameters_valid = gen_random_data_by_shape_and_type(l, r, shape_Parameters_valid, np.float16)
            C3_valid = gen_random_data_by_shape_and_type(l, r, shape_C3_valid, np.float16)

            print("BatchNum %d validABC generation finished."%(batchNum))

            temp = np.dot(Q_valid.astype(np.float32), KT_valid.astype(np.float32)).astype(np.float16)
            expect_result1_valid = alpha * temp + beta * C1_valid
            temp = np.dot(expect_result1_valid.astype(np.float32), V_valid.astype(np.float32)).astype(np.float16)
            expect_result2_valid = alpha * temp + beta * C2_valid
            temp = np.dot(expect_result2_valid.astype(np.float32), Parameters_valid.astype(np.float32)).astype(np.float16)
            expect_result3_valid = alpha * temp + beta * C3_valid

            print("BatchNum %d validExpectResult generation finished." % (batchNum))

            if check:
                print("expect_result1_valid: ")
                print(expect_result1_valid)
                print("expect_result2_valid: ")
                print(expect_result2_valid)
                print("expect_result3_valid: ")
                print(expect_result3_valid)
            

            maskA_zero_padding = np.zeros((1, wholeM), dtype=np.float16)
            
            Q_zero_padding = np.zeros((wholeM, wholeK), dtype=np.float16)
            KT_zero_padding = np.zeros((wholeK, wholeM), dtype=np.float16)
            C1_zero_padding = np.zeros((wholeM, wholeM), dtype=np.float16)
            expect_result1_zero_padding = np.zeros((wholeM, wholeM), dtype=np.float16)

            V_zero_padding = np.zeros((wholeM, wholeK), dtype=np.float16)
            C2_zero_padding = np.zeros((wholeM, wholeK), dtype=np.float16)
            expect_result2_zero_padding = np.zeros((wholeM, wholeK), dtype=np.float16)

            Parameters_zero_padding = np.zeros((wholeK, wholeN), dtype=np.float16)
            C3_zero_padding = np.zeros((wholeM, wholeN), dtype=np.float16)
            expect_result3_zero_padding = np.zeros((wholeM, wholeN), dtype=np.float16)


            maskA_zero_padding[:1, :validM] = maskA_valid

            Q_zero_padding[:validM, :wholeK] = Q_valid
            KT_zero_padding[:wholeK, :validM] = KT_valid
            C1_zero_padding[:validM, :validM] = C1_valid
            expect_result1_zero_padding[:validM, :validM] = expect_result1_valid

            V_zero_padding[:validM, :wholeK] = V_valid
            C2_zero_padding[:validM, :wholeK] = C2_valid
            expect_result2_zero_padding[:validM, :wholeK] = expect_result2_valid

            Parameters_zero_padding[:wholeK, :wholeN] = Parameters_valid
            C3_zero_padding[:validM, :wholeN] = C3_valid
            expect_result3_zero_padding[:validM, :wholeN] = expect_result3_valid

            print("BatchNum %d padding data generation finished." % (batchNum))

            if check:
                print("maskA_zero_padding: ")
                print(maskA_zero_padding)
                print("Q_zero_padding: ")
                print(Q_zero_padding)
                print("KT_zero_padding: ")
                print(KT_zero_padding)
                print("expect_result1_zero_padding: ")
                print(expect_result1_zero_padding)
                print("expect_result2_zero_padding: ")
                print(expect_result2_zero_padding)
                print("expect_result3_zero_padding: ")
                print(expect_result3_zero_padding)

            # tofile
            Q_zero_padding.tofile(Q_file, sep=" " if check else "")
            KT_zero_padding.tofile(KT_file, sep=" " if check else "")
            C1_zero_padding.tofile(C1_file, sep=" " if check else "")
            expect_result1_zero_padding.tofile(expect_result1_file, sep=" " if check else "")

            V_zero_padding.tofile(V_file, sep=" " if check else "")
            C2_zero_padding.tofile(C2_file, sep=" " if check else "")
            expect_result2_zero_padding.tofile(expect_result2_file, sep=" " if check else "")

            Parameters_zero_padding.tofile(Parameters_file, sep=" " if check else "")
            C3_zero_padding.tofile(C3_file, sep=" " if check else "")
            expect_result3_zero_padding.tofile(expect_result3_file, sep=" " if check else "")

            maskA_zero_padding.tofile(maskA_file, sep=" " if check else "")
            validM.tofile(validM_file, sep=" " if check else "")
            
            print("BatchNum %d tofile finished."%(batchNum))
            print("="*10)
            pass  # for batch

        pass  # with open

    pass  # gen_data_cal_expect_write_bin


if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument("wholeM", action="store", type=int)
    parser.add_argument("wholeN", action="store", type=int)
    parser.add_argument("wholeK", action="store", type=int)
    parser.add_argument("batchCount", action="store", type=int)

    args = parser.parse_args()

    wholeM = args.wholeM
    wholeN = args.wholeN
    wholeK = args.wholeK
    batchCount = args.batchCount

    gen_data_cal_expect_write_bin(wholeM, wholeN, wholeK, batchCount, False)